package cn.bitkit.component.common;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

/**
 * @author admin
 */
public class RateLimiter {

    private final Semaphore semaphore;
    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

    public RateLimiter(int qps) {
        this.semaphore = new Semaphore(qps);
        // 每秒释放一次许可
        scheduler.scheduleAtFixedRate(() -> {
            int permits = qps - semaphore.availablePermits();
            if (permits > 0) {
                semaphore.release(permits);
            }
        }, 0, 1, TimeUnit.SECONDS);
    }

    public void acquire() throws InterruptedException {
        semaphore.acquire();
    }

    public void shutdown() {
        scheduler.shutdown();
    }
}
